from langchain_ollama import OllamaEmbeddings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# 配置ollama 服务
ollama_url = "http://127.0.0.1:11434"


def test_embedding_consistency():
    ### 可切换不同的向量嵌入模型
    embed_model = OllamaEmbeddings(base_url=ollama_url, model="nomic-embed-text:latest")

    text1 = "人工智能技术"
    text2 = "人工智能技术"  # 完全相同
    text3 = "人工 智能 技术"  # 稍有不同
    text4 = "机器学习算法"  # 完全不同

    print("=== 向量一致性测试 ===")

    texts = [text1, text2, text3, text4]
    embeddings = []

    for text in texts:

        embedding = embed_model.embed_query(text)
        #embedding = embed_model.get_text_embedding(text) # 如果是hugging face 的运行的嵌入模型 使用这个
        embeddings.append(embedding)
        print(f"文本: '{text}'")
        print(f"向量前3维: {embedding[:3]}")

    # 计算相似度
    def cosine_similarity(a, b):
        import numpy as np
        return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

    print(f"\n相似度分析:")
    print(f"完全相同文本: {cosine_similarity(embeddings[0], embeddings[1]):.6f}")
    print(f"稍不同文本: {cosine_similarity(embeddings[0], embeddings[2]):.6f}")
    print(f"完全不同文本: {cosine_similarity(embeddings[0], embeddings[3]):.6f}")


test_embedding_consistency()
